{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# default_exp trick_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# trick_model\n",
    "\n",
    "> Model with ResNet Trick for experiments."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# setup and imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install git+https://github.com/ayasyrev/model_constructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pip install git+https://github.com/kornia/kornia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from kornia.contrib import MaxBlurPool2d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "# hide\n",
    "from nbdev.showdoc import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "from fastai.basic_train import *\n",
    "from fastai.vision import *\n",
    "from fastai.script import *\n",
    "from model_constructor.net import *\n",
    "from model_constructor.layers import SimpleSelfAttention, ConvLayer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#export\n",
    "import math\n",
    "import torch\n",
    "from torch.optim.optimizer import Optimizer, required\n",
    "import itertools as it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "from imagenette_experiments.train_utils import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# New ResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "class NewResBlock(Module):\n",
    "    def __init__(self, expansion, ni, nh, stride=1, \n",
    "                 conv_layer=ConvLayer, act_fn=act_fn,\n",
    "                 pool=nn.AvgPool2d(2, ceil_mode=True), sa=False,sym=False, zero_bn=True):\n",
    "        nf,ni = nh*expansion,ni*expansion\n",
    "        self.reduce = noop if stride==1 else pool\n",
    "        layers  = [(f\"conv_0\", conv_layer(ni, nh, 3, stride=stride, act_fn=act_fn)),\n",
    "                   (f\"conv_1\", conv_layer(nh, nf, 3, zero_bn=zero_bn, act=False))\n",
    "        ] if expansion == 1 else [\n",
    "                   (f\"conv_0\",conv_layer(ni, nh, 1, act_fn=act_fn)),\n",
    "                   (f\"conv_1\",conv_layer(nh, nh, 3, stride=1, act_fn=act_fn)), #!!!\n",
    "                   (f\"conv_2\",conv_layer(nh, nf, 1, zero_bn=zero_bn, act=False))\n",
    "        ]\n",
    "        if sa: layers.append(('sa', SimpleSelfAttention(nf,ks=1,sym=sym)))\n",
    "        self.convs = nn.Sequential(OrderedDict(layers))\n",
    "        self.idconv = noop if ni==nf else conv_layer(ni, nf, 1, act=False)\n",
    "        self.merge =act_fn\n",
    "\n",
    "    def forward(self, x): \n",
    "        o = self.reduce(x)\n",
    "        return self.merge(self.convs(o) + self.idconv(o))\n",
    "'''"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "NewResBlock now defined at model-constructor, so we alredy has it as we import Net from there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "model_constructor.net.NewResBlock"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "NewResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show_doc(NewResBlock)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model Constructor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# export\n",
    "class Model(Net):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.name = 'xresnet50_trick'\n",
    "        self.c_out = 10\n",
    "        self.expansion=4\n",
    "        self.layers=[3, 4,  6, 3]\n",
    "        self.stem_sizes = [3,32,64,64]\n",
    "        self.act_fn= Mish()\n",
    "        self.sa = True\n",
    "        self.pool = MaxBlurPool2d(3, True)\n",
    "        self.stem_pool = self.pool\n",
    "        self.block = NewResBlock"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       " constr xresnet50_trick"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## repr model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  model xresnet50_trick\n",
       "  (stem): Sequential(\n",
       "    (conv_0): ConvLayer(\n",
       "      (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): Mish()\n",
       "    )\n",
       "    (conv_1): ConvLayer(\n",
       "      (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): Mish()\n",
       "    )\n",
       "    (conv_2): ConvLayer(\n",
       "      (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (act_fn): Mish()\n",
       "    )\n",
       "    (stem_pool): MaxBlurPool2d()\n",
       "  )\n",
       "  (body): Sequential(\n",
       "    (l_0): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): ConvLayer(\n",
       "          (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_2): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "          (sa): SimpleSelfAttention(\n",
       "            (conv): Conv1d(256, 256, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "    )\n",
       "    (l_1): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (reduce): MaxBlurPool2d()\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): ConvLayer(\n",
       "          (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_2): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_3): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "    )\n",
       "    (l_2): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (reduce): MaxBlurPool2d()\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): ConvLayer(\n",
       "          (conv): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_2): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_3): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_4): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_5): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "    )\n",
       "    (l_3): Sequential(\n",
       "      (bl_0): NewResBlock(\n",
       "        (reduce): MaxBlurPool2d()\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (idconv): ConvLayer(\n",
       "          (conv): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_1): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "      (bl_2): NewResBlock(\n",
       "        (convs): Sequential(\n",
       "          (conv_0): ConvLayer(\n",
       "            (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_1): ConvLayer(\n",
       "            (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "            (act_fn): Mish()\n",
       "          )\n",
       "          (conv_2): ConvLayer(\n",
       "            (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "            (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          )\n",
       "        )\n",
       "        (merge): Mish()\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (head): Sequential(\n",
       "    (pool): AdaptiveAvgPool2d(output_size=1)\n",
       "    (flat): Flatten()\n",
       "    (fc): Linear(in_features=2048, out_features=10, bias=True)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (conv_0): ConvLayer(\n",
       "    (conv): Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): Mish()\n",
       "  )\n",
       "  (conv_1): ConvLayer(\n",
       "    (conv): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): Mish()\n",
       "  )\n",
       "  (conv_2): ConvLayer(\n",
       "    (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (act_fn): Mish()\n",
       "  )\n",
       "  (stem_pool): MaxBlurPool2d()\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.stem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (l_0): Sequential(\n",
       "    (bl_0): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (idconv): ConvLayer(\n",
       "        (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_1): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_2): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (sa): SimpleSelfAttention(\n",
       "          (conv): Conv1d(256, 256, kernel_size=(1,), stride=(1,), bias=False)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "  )\n",
       "  (l_1): Sequential(\n",
       "    (bl_0): NewResBlock(\n",
       "      (reduce): MaxBlurPool2d()\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (idconv): ConvLayer(\n",
       "        (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_1): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_2): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_3): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "  )\n",
       "  (l_2): Sequential(\n",
       "    (bl_0): NewResBlock(\n",
       "      (reduce): MaxBlurPool2d()\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (idconv): ConvLayer(\n",
       "        (conv): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_1): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_2): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_3): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_4): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_5): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "  )\n",
       "  (l_3): Sequential(\n",
       "    (bl_0): NewResBlock(\n",
       "      (reduce): MaxBlurPool2d()\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (idconv): ConvLayer(\n",
       "        (conv): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_1): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "    (bl_2): NewResBlock(\n",
       "      (convs): Sequential(\n",
       "        (conv_0): ConvLayer(\n",
       "          (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_1): ConvLayer(\n",
       "          (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act_fn): Mish()\n",
       "        )\n",
       "        (conv_2): ConvLayer(\n",
       "          (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "          (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (merge): Mish()\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.body"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (pool): AdaptiveAvgPool2d(output_size=1)\n",
       "  (flat): Flatten()\n",
       "  (fc): Linear(in_features=2048, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (bl_0): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_1): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_2): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (sa): SimpleSelfAttention(\n",
       "        (conv): Conv1d(256, 256, kernel_size=(1,), stride=(1,), bias=False)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.body.l_0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (bl_0): NewResBlock(\n",
       "    (reduce): MaxBlurPool2d()\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(256, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_1): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_2): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_3): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.body.l_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (bl_0): NewResBlock(\n",
       "    (reduce): MaxBlurPool2d()\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(512, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_1): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_2): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_3): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_4): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_5): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.body.l_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (bl_0): NewResBlock(\n",
       "    (reduce): MaxBlurPool2d()\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (idconv): ConvLayer(\n",
       "      (conv): Conv2d(1024, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "      (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_1): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       "  (bl_2): NewResBlock(\n",
       "    (convs): Sequential(\n",
       "      (conv_0): ConvLayer(\n",
       "        (conv): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_1): ConvLayer(\n",
       "        (conv): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        (act_fn): Mish()\n",
       "      )\n",
       "      (conv_2): ConvLayer(\n",
       "        (conv): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)\n",
       "        (bn): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "    )\n",
       "    (merge): Mish()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# hide\n",
    "model.body.l_3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
